from langchain_community.chat_models import ChatTongyi
from langchain_core.prompts.prompt import PromptTemplate
from langchain.prompts.few_shot import FewShotPromptTemplate
from langchain.prompts.example_selector import SemanticSimilarityExampleSelector
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Chroma

import os

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"
llm = ChatTongyi()

# 小样本示例
examples = [
    {
        "question": "你是谁?",
        "answer":
            """
            我是wenwenc9
            """
    },
    {
        "question": "爸爸叫什么",
        "answer":
            """
            叫爷爷
            """
    },
    {
        "question": "朱自清是谁",
        "answer":
            """
            朱自清是文学家
            """
    }
]

# 用于格式化输入变量（在这里是"question"和"answer"）到具体的字符串中。这个模板的结构是：“Question: ”后面跟着问题，然后换行，再跟上回答
example_prompt = PromptTemplate(input_variables=["question", "answer"], template="Question: {question}\n{answer}")
prompt = FewShotPromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
    suffix="Question: {input}",  # 后缀
    input_variables=["input"]
)
# 选择器
example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,
    DashScopeEmbeddings(),
    Chroma,
    k=1)
# 创建一个问题
question = "朱自清"

# 根据向量引擎从中选择一个符号
selected_examples = example_selector.select_examples({"input": question})
print(selected_examples)
